﻿// Inner Fire 游戏引擎库
// 根签名管理类
//
// Copyright (c) 2024 by 尔西大帝. All rights reserved.
//
// 这个类提供根签名的模块化管理。
//
// Author: 尔西大帝
// Email: 2670613871@qq.com
// Created: 2024-12-07

#include "dx_root_signature.h"
#include "dx_exception.h"
#include <array>
#include <klog.h>
#include <dx/graphics/dx_graphics.h>

namespace ifire::dx {
// 获得标准RootSignature
std::array<const CD3DX12_STATIC_SAMPLER_DESC, 7> GetStandardStaticSamplers();

DxRootSignature::DxRootSignature(DxGraphics* dx) : DxModule(dx) {}

void DxRootSignature::BuildRootSignature() {
  BuildStandard();
  BuildSsao();
  BuildPostProcess();
  BuildWaves();
  BuildUI();
  BuildParticle();
}

void DxRootSignature::BuildStandard() {
  const int SKYBOX_REGISTER = 0;
  const int TEXTURE_REGISTER = 3;
  const int NORMAL_REGISTER = 4;
  const int TEXTURE_ARRAY_REGISTER = 5; // = t5
  const int DISPLACEMENT_REGISTER = 3;  // = t6

  // 用于纹理的描述符表
  CD3DX12_DESCRIPTOR_RANGE texTable;
  texTable.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, TEXTURE_REGISTER);
  // 用于纹理数组的描述符表
  CD3DX12_DESCRIPTOR_RANGE tex_arr_table;
  tex_arr_table.Init(
      D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, TEXTURE_ARRAY_REGISTER);

  // 用于法线贴图的描述符表
  CD3DX12_DESCRIPTOR_RANGE normal_table;
  normal_table.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, NORMAL_REGISTER);

  // 用于天空盒贴图的描述符表。这是连续3个的定义，包括天空盒，阴影，SSAO。
  // 所以在描述符堆中也需要一样连续。
  CD3DX12_DESCRIPTOR_RANGE cube_table;
  cube_table.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 3, SKYBOX_REGISTER);

  // 用于置换贴图的描述符表
  CD3DX12_DESCRIPTOR_RANGE displacement_table;
  displacement_table.Init(
      D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, DISPLACEMENT_REGISTER, 1);

  // 根参数可以是描述符表，根描述符，或根常量
  CD3DX12_ROOT_PARAMETER params[ShaderSlot::SLOTS];
  // 设置描述符区域的数量，指向描述符区域数组的指针
  params[ShaderSlot::TEXTURE].InitAsDescriptorTable(
      1, &texTable, D3D12_SHADER_VISIBILITY_PIXEL);
  params[ShaderSlot::OBJECT_CB].InitAsConstantBufferView(0);
  params[ShaderSlot::MAIN_PASS].InitAsConstantBufferView(1);
  params[ShaderSlot::SKINNED_CB].InitAsConstantBufferView(2);
  // 这个东西用于存材质
  params[ShaderSlot::MATERIAL].InitAsShaderResourceView(0, 1);
  // 用于法线贴图
  params[ShaderSlot::NORMAL_MAP].InitAsDescriptorTable(
      1, &normal_table, D3D12_SHADER_VISIBILITY_PIXEL);
  // 用于天空盒贴图
  params[ShaderSlot::SKYBOX].InitAsDescriptorTable(
      1, &cube_table, D3D12_SHADER_VISIBILITY_PIXEL);
  // 用于纹理数组
  params[ShaderSlot::TEXTURE_ARRAY].InitAsDescriptorTable(
      1, &tex_arr_table, D3D12_SHADER_VISIBILITY_PIXEL);
  // 用于置换贴图
  // 注意这里我被教程坑了一把，这里应该是D3D12_SHADER_VISIBILITY_ALL，
  // 表示所有着色器都能访问。不然的话，在Vertex Shader中就会出 is not fully
  // bound in root signature 错误。
  params[ShaderSlot::DISPLACEMENT].InitAsDescriptorTable(
      1, &displacement_table, D3D12_SHADER_VISIBILITY_ALL);

  // 注意这里和之前的变化，增加了采样器的数量和地址。
  auto samplers = GetStandardStaticSamplers();

  // 根签名由一组根参数构成
  CD3DX12_ROOT_SIGNATURE_DESC desc(ShaderSlot::SLOTS, params,
      static_cast<UINT>(samplers.size()), samplers.data(),
      D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

  // 创建签名
  CreateRootSignature(desc, mStandard);
}

void DxRootSignature::CreateRootSignature(
    const CD3DX12_ROOT_SIGNATURE_DESC& desc,
    ComPtr<ID3D12RootSignature>& sign) {
  // 序列化根签名
  ComPtr<ID3DBlob> serialized_root_sign = nullptr;
  ComPtr<ID3DBlob> error_blob = nullptr;
  HRESULT hr = D3D12SerializeRootSignature(&desc, D3D_ROOT_SIGNATURE_VERSION_1,
      serialized_root_sign.GetAddressOf(), error_blob.GetAddressOf());
  // 常规错误检查
  if (error_blob != nullptr) {
    KERROR((char*)error_blob->GetBufferPointer());
  }
  HASSERT(hr);

  // 创建根签名
  HASSERT(dx_->Device()->CreateRootSignature(0, //
      serialized_root_sign->GetBufferPointer(),
      serialized_root_sign->GetBufferSize(),
      IID_PPV_ARGS(sign.GetAddressOf())));
}

///////////////////////////////////////////////////////////////////////////////
// 默认采样器
///////////////////////////////////////////////////////////////////////////////
const CD3DX12_STATIC_SAMPLER_DESC pointWrap(0, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_POINT,            // filter
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,           // addressU
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,           // addressV
    D3D12_TEXTURE_ADDRESS_MODE_WRAP);          // addressW

const CD3DX12_STATIC_SAMPLER_DESC pointClamp(1, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_POINT,             // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,           // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,           // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP);          // addressW

const CD3DX12_STATIC_SAMPLER_DESC linearWrap(2, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_LINEAR,            // filter
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,            // addressU
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,            // addressV
    D3D12_TEXTURE_ADDRESS_MODE_WRAP);           // addressW

const CD3DX12_STATIC_SAMPLER_DESC linearClamp(3, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_LINEAR,             // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,            // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,            // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP);           // addressW

const CD3DX12_STATIC_SAMPLER_DESC anisotropicWrap(4, // shaderRegister
    D3D12_FILTER_ANISOTROPIC,                        // filter
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,                 // addressU
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,                 // addressV
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,                 // addressW
    0.0f,                                            // mipLODBias
    8);                                              // maxAnisotropy

const CD3DX12_STATIC_SAMPLER_DESC anisotropicClamp(5, // shaderRegister
    D3D12_FILTER_ANISOTROPIC,                         // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressW
    0.0f,                                             // mipLODBias
    8);                                               // maxAnisotropy

const CD3DX12_STATIC_SAMPLER_DESC shadowSampler(6,    // shaderRegister
    D3D12_FILTER_COMPARISON_MIN_MAG_LINEAR_MIP_POINT, // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressW
    0.0f,                                             // mipLODBias
    16,                                               // maxAnisotropy
    D3D12_COMPARISON_FUNC_LESS_EQUAL, D3D12_STATIC_BORDER_COLOR_OPAQUE_BLACK);

const CD3DX12_STATIC_SAMPLER_DESC SSAO_POINT_CLAMP(0, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_POINT,                   // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                 // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP);                // addressW

const CD3DX12_STATIC_SAMPLER_DESC SSAO_LINEAR_CLAMP(1, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_LINEAR,                   // filter
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                  // addressU
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP,                  // addressV
    D3D12_TEXTURE_ADDRESS_MODE_CLAMP);                 // addressW

const CD3DX12_STATIC_SAMPLER_DESC SSAO_DEPTH_MAP(2, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_LINEAR,                // filter
    D3D12_TEXTURE_ADDRESS_MODE_BORDER,              // addressU
    D3D12_TEXTURE_ADDRESS_MODE_BORDER,              // addressV
    D3D12_TEXTURE_ADDRESS_MODE_BORDER,              // addressW
    0.0f, 0, D3D12_COMPARISON_FUNC_LESS_EQUAL,
    D3D12_STATIC_BORDER_COLOR_OPAQUE_WHITE);

const CD3DX12_STATIC_SAMPLER_DESC SSAO_LINEAR_WRAP(3, // shaderRegister
    D3D12_FILTER_MIN_MAG_MIP_LINEAR,                  // filter
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,                  // addressU
    D3D12_TEXTURE_ADDRESS_MODE_WRAP,                  // addressV
    D3D12_TEXTURE_ADDRESS_MODE_WRAP);                 // addressW

std::array<const CD3DX12_STATIC_SAMPLER_DESC, 7> GetStandardStaticSamplers() {
  // 应用程序通常只需要少量的采样器。因此，只需提前定义它们，并将其作为根签名的一部分保持可用。
  return {pointWrap, pointClamp, linearWrap, linearClamp, anisotropicWrap,
      anisotropicClamp, shadowSampler};
}

void DxRootSignature::BuildSsao() {
  CD3DX12_DESCRIPTOR_RANGE texTable0;
  texTable0.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 2, 0, 0);

  CD3DX12_DESCRIPTOR_RANGE texTable1;
  texTable1.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, 2, 0);

  // Root parameter can be a table, root descriptor or root constants.
  CD3DX12_ROOT_PARAMETER params[4];

  // Perfomance TIP: Order from most frequent to least frequent.
  params[0].InitAsConstantBufferView(0);
  params[1].InitAsConstants(1, 1);
  params[2].InitAsDescriptorTable(1, &texTable0, D3D12_SHADER_VISIBILITY_PIXEL);
  params[3].InitAsDescriptorTable(1, &texTable1, D3D12_SHADER_VISIBILITY_PIXEL);

  std::array<CD3DX12_STATIC_SAMPLER_DESC, 4> staticSamplers = {
      SSAO_POINT_CLAMP, SSAO_LINEAR_CLAMP, SSAO_DEPTH_MAP, SSAO_LINEAR_WRAP};

  // A root signature is an array of root parameters.
  CD3DX12_ROOT_SIGNATURE_DESC desc(4, params, (UINT)staticSamplers.size(),
      staticSamplers.data(),
      D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

  CreateRootSignature(desc, mSsao);
}

void DxRootSignature::BuildWaves() {
  CD3DX12_DESCRIPTOR_RANGE uavTable0;
  uavTable0.Init(D3D12_DESCRIPTOR_RANGE_TYPE_UAV, 1, 0);

  CD3DX12_DESCRIPTOR_RANGE uavTable1;
  uavTable1.Init(D3D12_DESCRIPTOR_RANGE_TYPE_UAV, 1, 1);

  CD3DX12_DESCRIPTOR_RANGE uavTable2;
  uavTable2.Init(D3D12_DESCRIPTOR_RANGE_TYPE_UAV, 1, 2);

  // Root parameter can be a table, root descriptor or root constants.
  CD3DX12_ROOT_PARAMETER params[4];

  // Perfomance TIP: Order from most frequent to least frequent.
  params[0].InitAsConstants(6, 0);
  params[1].InitAsDescriptorTable(1, &uavTable0);
  params[2].InitAsDescriptorTable(1, &uavTable1);
  params[3].InitAsDescriptorTable(1, &uavTable2);

  // A root signature is an array of root parameters.
  CD3DX12_ROOT_SIGNATURE_DESC desc(
      4, params, 0, nullptr, D3D12_ROOT_SIGNATURE_FLAG_NONE);

  CreateRootSignature(desc, mWaves);
}

void DxRootSignature::BuildPostProcess() {
  CD3DX12_DESCRIPTOR_RANGE srvTable0;
  srvTable0.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, 0);

  CD3DX12_DESCRIPTOR_RANGE srvTable1;
  srvTable1.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, 1);

  CD3DX12_DESCRIPTOR_RANGE uavTable0;
  uavTable0.Init(D3D12_DESCRIPTOR_RANGE_TYPE_UAV, 1, 0);

  // Root parameter can be a table, root descriptor or root constants.
  CD3DX12_ROOT_PARAMETER params[3];

  // Perfomance TIP: Order from most frequent to least frequent.
  params[0].InitAsDescriptorTable(1, &srvTable0);
  params[1].InitAsDescriptorTable(1, &srvTable1);
  params[2].InitAsDescriptorTable(1, &uavTable0);

  auto staticSamplers = GetStandardStaticSamplers();

  // A root signature is an array of root parameters.
  CD3DX12_ROOT_SIGNATURE_DESC desc(3, params, (UINT)staticSamplers.size(),
      staticSamplers.data(),
      D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

  CreateRootSignature(desc, mPostProcessing);
}

void DxRootSignature::BuildUI() {
  const auto SLOT_COUNT = 2;

  // 用于UI的图片
  CD3DX12_DESCRIPTOR_RANGE texTable;
  texTable.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, 0);

  // Root parameter can be a table, root descriptor or root constants.
  CD3DX12_ROOT_PARAMETER params[SLOT_COUNT];

  // Perfomance TIP: Order from most frequent to least frequent.
  params[0].InitAsDescriptorTable(1, &texTable, D3D12_SHADER_VISIBILITY_PIXEL);
  // 用于实例化的项目，在t1,space1中。
  params[1].InitAsShaderResourceView(0, 1);

  std::array<CD3DX12_STATIC_SAMPLER_DESC, 6> staticSamplers = {pointWrap,
      pointClamp, linearWrap, linearClamp, anisotropicWrap, anisotropicClamp};

  // A root signature is an array of root parameters.
  CD3DX12_ROOT_SIGNATURE_DESC desc(SLOT_COUNT, params,
      (UINT)staticSamplers.size(), staticSamplers.data(),
      D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

  CreateRootSignature(desc, mUI);
}

void DxRootSignature::BuildParticle() {
  const auto SLOT_COUNT = 4;

  // 用于UI的图片
  CD3DX12_DESCRIPTOR_RANGE texTable;
  texTable.Init(D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 1, 0);

  // Root parameter can be a table, root descriptor or root constants.
  CD3DX12_ROOT_PARAMETER params[SLOT_COUNT];

  // 用于指定贴图。这里使用RootSignatureSlot::系列的索引，是因为能够保持顺序一致就保持一一致。
  params[ShaderSlot::TEXTURE].InitAsDescriptorTable(
      1, &texTable, D3D12_SHADER_VISIBILITY_PIXEL);
  params[ShaderSlot::OBJECT_CB].InitAsConstantBufferView(0);
  params[ShaderSlot::MAIN_PASS].InitAsConstantBufferView(1);
  // 用于实例化的项目，在t1,space1中。
  params[ShaderSlot::PARTICLE_DATA].InitAsShaderResourceView(0, 1);

  std::array<CD3DX12_STATIC_SAMPLER_DESC, 6> staticSamplers = {pointWrap,
      pointClamp, linearWrap, linearClamp, anisotropicWrap, anisotropicClamp};

  // A root signature is an array of root parameters.
  CD3DX12_ROOT_SIGNATURE_DESC desc(SLOT_COUNT, params,
      (UINT)staticSamplers.size(), staticSamplers.data(),
      D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT);

  CreateRootSignature(desc, mParticle);
}

} // namespace ifire::dx